{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "38e73e84-12d9-42c0-8c97-62a358900bcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import os.path as osp\n",
    "from torch.utils import data as data\n",
    "from torchvision.transforms.functional import normalize\n",
    "\n",
    "from basicsr.data.data_util import paths_from_lmdb, scandir\n",
    "from basicsr.data.transforms import augment, paired_random_crop\n",
    "from basicsr.utils import FileClient, imfrombytes, img2tensor\n",
    "from basicsr.utils.matlab_functions import imresize, rgb2ycbcr\n",
    "from basicsr.utils.registry import DATASET_REGISTRY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bf6bb909-87d7-4be2-b81a-d8ce49f502bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.nn import functional as F\n",
    "\n",
    "from basicsr.utils.registry import MODEL_REGISTRY\n",
    "from basicsr.models.sr_model import SRModel\n",
    "from basicsr.metrics import calculate_metric\n",
    "from basicsr.utils import imwrite, tensor2img\n",
    "\n",
    "import math\n",
    "from tqdm import tqdm\n",
    "from os import path as osp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fcb2af44-79ee-445a-861d-61a68d470217",
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "file=\"train_HAT_SRx2_from_scratch.yml\"\n",
    "with open(file,'r') as file:\n",
    "    opt=yaml.safe_load(file)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e768196-86d1-4f8d-a861-2f389d1c7bfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.checkpoint as checkpoint\n",
    "\n",
    "from basicsr.utils.registry import ARCH_REGISTRY\n",
    "from basicsr.archs.arch_util import to_2tuple, trunc_normal_\n",
    "\n",
    "from einops import rearrange\n",
    "\n",
    "def drop_path(x, drop_prob: float = 0., training: bool = False):\n",
    "    \"\"\"Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).\n",
    "\n",
    "    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py\n",
    "    \"\"\"\n",
    "    if drop_prob == 0. or not training:\n",
    "        return x\n",
    "    keep_prob = 1 - drop_prob\n",
    "    shape = (x.shape[0], ) + (1, ) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets\n",
    "    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)\n",
    "    random_tensor.floor_()  # binarize\n",
    "    output = x.div(keep_prob) * random_tensor\n",
    "    return output\n",
    "\n",
    "\n",
    "class DropPath(nn.Module):\n",
    "    \"\"\"Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).\n",
    "\n",
    "    From: https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, drop_prob=None):\n",
    "        super(DropPath, self).__init__()\n",
    "        self.drop_prob = drop_prob\n",
    "\n",
    "    def forward(self, x):\n",
    "        return drop_path(x, self.drop_prob, self.training)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "7d7b4a34-7ed8-4a61-8391-45e787ce6dd1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 3, 10, 27])"
      ]
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x=torch.tensor([1,2,3])\n",
    "y=torch.tensor([3,5,9])\n",
    "x*y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c125e3c2-1fcf-404d-86bd-9ed22e92f008",
   "metadata": {},
   "outputs": [],
   "source": [
    "#实现通道加权\n",
    "class ChannelAttention(nn.Module):     \n",
    "    \"\"\"Channel attention used in RCAN.\n",
    "    Args:\n",
    "        num_feat (int): Channel number of intermediate features.\n",
    "        squeeze_factor (int): Channel squeeze factor. Default: 16.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, num_feat, squeeze_factor=16):\n",
    "        super(ChannelAttention, self).__init__()\n",
    "        self.attention = nn.Sequential(\n",
    "            nn.AdaptiveAvgPool2d(1),\n",
    "            nn.Conv2d(num_feat, num_feat // squeeze_factor, 1, padding=0),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Conv2d(num_feat // squeeze_factor, num_feat, 1, padding=0),\n",
    "            nn.Sigmoid())\n",
    "\n",
    "    def forward(self, x):\n",
    "        y = self.attention(x)\n",
    "        return x * y\n",
    "\n",
    "class CAB(nn.Module):\n",
    "\n",
    "    def __init__(self, num_feat, compress_ratio=3, squeeze_factor=30):\n",
    "        super(CAB, self).__init__()\n",
    "\n",
    "        self.cab = nn.Sequential(\n",
    "            nn.Conv2d(num_feat, num_feat // compress_ratio, 3, 1, 1),\n",
    "            nn.GELU(),\n",
    "            nn.Conv2d(num_feat // compress_ratio, num_feat, 3, 1, 1),\n",
    "            ChannelAttention(num_feat, squeeze_factor)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.cab(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2ecd2e4-b8e2-4e54-b3b5-11bf0c7bce28",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mlp(nn.Module):\n",
    "\n",
    "    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n",
    "        super().__init__()\n",
    "        out_features = out_features or in_features\n",
    "        hidden_features = hidden_features or in_features\n",
    "        self.fc1 = nn.Linear(in_features, hidden_features)\n",
    "        self.act = act_layer()\n",
    "        self.fc2 = nn.Linear(hidden_features, out_features)\n",
    "        self.drop = nn.Dropout(drop)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.fc1(x)\n",
    "        x = self.act(x)\n",
    "        x = self.drop(x)\n",
    "        x = self.fc2(x)\n",
    "        x = self.drop(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b7edae7-5533-444e-a722-390a47a0d704",
   "metadata": {},
   "outputs": [],
   "source": [
    "#将大的图像分割为多个小的图像\n",
    "def window_partition(x, window_size):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        x: (b, h, w, c)\n",
    "        window_size (int): window size\n",
    "\n",
    "    Returns:\n",
    "        windows: (num_windows*b, window_size, window_size, c)\n",
    "    \"\"\"\n",
    "    b, h, w, c = x.shape\n",
    "    x = x.view(b, h // window_size, window_size, w // window_size, window_size, c)\n",
    "    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, c)\n",
    "    return windows\n",
    "\n",
    "#和并图像\n",
    "def window_reverse(windows, window_size, h, w):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        windows: (num_windows*b, window_size, window_size, c)\n",
    "        window_size (int): Window size\n",
    "        h (int): Height of image\n",
    "        w (int): Width of image\n",
    "\n",
    "    Returns:\n",
    "        x: (b, h, w, c)\n",
    "    \"\"\"\n",
    "    b = int(windows.shape[0] / (h * w / window_size / window_size))\n",
    "    x = windows.view(b, h // window_size, w // window_size, window_size, window_size, -1)\n",
    "    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "c48a93c1-56c4-425f-9022-2196e7a30157",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WindowAttention(nn.Module):\n",
    "    r\"\"\" Window based multi-head self attention (W-MSA) module with relative position bias.\n",
    "    It supports both of shifted and non-shifted window.\n",
    "\n",
    "    Args:\n",
    "        dim (int): Number of input channels.\n",
    "        window_size (tuple[int]): The height and width of the window.\n",
    "        num_heads (int): Number of attention heads.\n",
    "        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set\n",
    "        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0\n",
    "        proj_drop (float, optional): Dropout ratio of output. Default: 0.0\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0.1, proj_drop=0.1):\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.window_size = window_size  # Wh, Ww\n",
    "        self.num_heads = num_heads\n",
    "        head_dim = dim // num_heads\n",
    "        #self.scale = qk_scale or head_dim**-0.5\n",
    "        self.scale =head_dim**-0.5\n",
    "\n",
    "        # define a parameter table of relative position bias\n",
    "        self.relative_position_bias_table = nn.Parameter(\n",
    "            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH\n",
    "\n",
    "        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n",
    "        self.attn_drop = nn.Dropout(attn_drop)\n",
    "        self.proj = nn.Linear(dim, dim)\n",
    "\n",
    "        self.proj_drop = nn.Dropout(proj_drop)\n",
    "\n",
    "        trunc_normal_(self.relative_position_bias_table, std=.02)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, x, rpi, mask=None):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            x: input features with shape of (num_windows*b, n, c)\n",
    "            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None\n",
    "        \"\"\"\n",
    "        b_, n, c = x.shape\n",
    "        qkv = self.qkv(x).reshape(b_, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4)\n",
    "        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)\n",
    "\n",
    "        q = q * self.scale\n",
    "        attn = (q @ k.transpose(-2, -1))\n",
    "\n",
    "        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(\n",
    "            self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH\n",
    "        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww\n",
    "        attn = attn + relative_position_bias.unsqueeze(0)\n",
    "\n",
    "        if mask is not None:\n",
    "            nw = mask.shape[0]\n",
    "            attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)\n",
    "            attn = attn.view(-1, self.num_heads, n, n)\n",
    "            attn = self.softmax(attn)\n",
    "        else:\n",
    "            attn = self.softmax(attn)\n",
    "\n",
    "        attn = self.attn_drop(attn)\n",
    "\n",
    "        x = (attn @ v).transpose(1, 2).reshape(b_, n, c)\n",
    "        x = self.proj(x)\n",
    "        x = self.proj_drop(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e7bdd2c-d9c1-44fd-9ee7-f227bfc06c3e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "d3baef68-0d4f-4698-9a9c-b8388af6cb44",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HAB(nn.Module):\n",
    "    r\"\"\" Hybrid Attention Block.\n",
    "\n",
    "    Args:\n",
    "        dim (int): Number of input channels.\n",
    "        input_resolution (tuple[int]): Input resolution.\n",
    "        num_heads (int): Number of attention heads.\n",
    "        window_size (int): Window size.\n",
    "        shift_size (int): Shift size for SW-MSA.\n",
    "        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n",
    "        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n",
    "        drop (float, optional): Dropout rate. Default: 0.0\n",
    "        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n",
    "        drop_path (float, optional): Stochastic depth rate. Default: 0.0\n",
    "        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU\n",
    "        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 dim,\n",
    "                 input_resolution,\n",
    "                 num_heads,\n",
    "                 window_size=8,\n",
    "                 shift_size=2,\n",
    "                 compress_ratio=3,\n",
    "                 squeeze_factor=30,\n",
    "                 conv_scale=0.01,\n",
    "                 mlp_ratio=4.,\n",
    "                 qkv_bias=True,\n",
    "                 qk_scale=None,\n",
    "                 drop=0.1,\n",
    "                 attn_drop=0.1,\n",
    "                 drop_path=0.1,\n",
    "                 act_layer=nn.GELU,\n",
    "                 norm_layer=nn.LayerNorm):\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.input_resolution = input_resolution\n",
    "        self.num_heads = num_heads\n",
    "        self.window_size = window_size\n",
    "        self.shift_size = shift_size\n",
    "        self.mlp_ratio = mlp_ratio\n",
    "        if min(self.input_resolution) <= self.window_size:\n",
    "            # if window size is larger than input resolution, we don't partition windows\n",
    "            self.shift_size = 0\n",
    "            self.window_size = min(self.input_resolution)\n",
    "        assert 0 <= self.shift_size < self.window_size, 'shift_size must in 0-window_size'\n",
    "\n",
    "        \n",
    "        self.norm1 = norm_layer(dim)\n",
    "        self.attn = WindowAttention(\n",
    "            dim,\n",
    "            window_size=to_2tuple(self.window_size),\n",
    "            num_heads=num_heads,\n",
    "            qkv_bias=qkv_bias,\n",
    "            qk_scale=qk_scale,\n",
    "            attn_drop=attn_drop,\n",
    "            proj_drop=drop)\n",
    "\n",
    "        self.conv_scale = conv_scale\n",
    "        self.conv_block = CAB(num_feat=dim, compress_ratio=compress_ratio, squeeze_factor=squeeze_factor)\n",
    "\n",
    "        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n",
    "        \n",
    "        self.norm2 = norm_layer(dim)\n",
    "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
    "        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n",
    "\n",
    "    def forward(self, x, x_size, rpi_sa, attn_mask):\n",
    "        h, w = x_size\n",
    "        b, _,c = x.shape\n",
    "        # assert seq_len == h * w, \"input feature has wrong size\"\n",
    "\n",
    "        shortcut = x\n",
    "\n",
    "        # self.norm1 =self.norm1((c,w,h))\n",
    "        x = self.norm1(x)\n",
    "        x = x.view(b, h, w, c)\n",
    "\n",
    "        # Conv_X\n",
    "        conv_x = self.conv_block(x.permute(0, 3, 1, 2))\n",
    "        conv_x = conv_x.permute(0, 2, 3, 1).contiguous().view(b, h * w, c)\n",
    "\n",
    "        # cyclic shift\n",
    "        if self.shift_size > 0:\n",
    "            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))\n",
    "            attn_mask = attn_mask\n",
    "        else:\n",
    "            shifted_x = x\n",
    "            attn_mask = None\n",
    "\n",
    "        # partition windows\n",
    "        x_windows = window_partition(shifted_x, self.window_size)  # nw*b, window_size, window_size, c\n",
    "        x_windows = x_windows.view(-1, self.window_size * self.window_size, c)  # nw*b, window_size*window_size, c\n",
    "\n",
    "        # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size\n",
    "        attn_windows = self.attn(x_windows, rpi=rpi_sa, mask=attn_mask)\n",
    "\n",
    "        # merge windows\n",
    "        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, c)\n",
    "        shifted_x = window_reverse(attn_windows, self.window_size, h, w)  # b h' w' c\n",
    "\n",
    "        # reverse cyclic shift\n",
    "        if self.shift_size > 0:\n",
    "            attn_x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))\n",
    "        else:\n",
    "            attn_x = shifted_x\n",
    "        attn_x = attn_x.view(b, h * w, c)\n",
    "\n",
    "        # FFN\n",
    "        x = shortcut + self.drop_path(attn_x) + conv_x * self.conv_scale\n",
    "        x = x + self.drop_path(self.mlp(self.norm2(x)))\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "4ff301f0-04bf-4d5b-9786-d1ee538b1cbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "hab=HAB(96,(256,256),8)\n",
    "# s=hab(torch.zeros(1,96,256,256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "31eb2e3f-c07b-4bed-b32c-9b0327e6f116",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "class PatchMerging(nn.Module):\n",
    "    r\"\"\" Patch Merging Layer.\n",
    "\n",
    "    Args:\n",
    "        input_resolution (tuple[int]): Resolution of input feature.\n",
    "        dim (int): Number of input channels.\n",
    "        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):\n",
    "        super().__init__()\n",
    "        self.input_resolution = input_resolution\n",
    "        self.dim = dim\n",
    "        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)\n",
    "        self.norm = norm_layer(4 * dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        x: b, h*w, c\n",
    "        \"\"\"\n",
    "        h, w = self.input_resolution\n",
    "        b, seq_len, c = x.shape\n",
    "        assert seq_len == h * w, 'input feature has wrong size'\n",
    "        assert h % 2 == 0 and w % 2 == 0, f'x size ({h}*{w}) are not even.'\n",
    "\n",
    "        x = x.view(b, h, w, c)\n",
    "\n",
    "        x0 = x[:, 0::2, 0::2, :]  # b h/2 w/2 c\n",
    "        x1 = x[:, 1::2, 0::2, :]  # b h/2 w/2 c\n",
    "        x2 = x[:, 0::2, 1::2, :]  # b h/2 w/2 c\n",
    "        x3 = x[:, 1::2, 1::2, :]  # b h/2 w/2 c\n",
    "        x = torch.cat([x0, x1, x2, x3], -1)  # b h/2 w/2 4*c\n",
    "        x = x.view(b, -1, 4 * c)  # b h/2*w/2 4*c\n",
    "\n",
    "        x = self.norm(x)\n",
    "        x = self.reduction(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "class OCAB(nn.Module):\n",
    "    # overlapping cross-attention block\n",
    "\n",
    "    def __init__(self, dim,\n",
    "                input_resolution,\n",
    "                window_size,\n",
    "                overlap_ratio,\n",
    "                num_heads,\n",
    "                qkv_bias=True,\n",
    "                qk_scale=None,\n",
    "                mlp_ratio=2,\n",
    "                norm_layer=nn.LayerNorm\n",
    "                ):\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.input_resolution = input_resolution\n",
    "        self.window_size = window_size\n",
    "        self.num_heads = num_heads\n",
    "        head_dim = dim // num_heads\n",
    "        self.scale = qk_scale or head_dim**-0.5\n",
    "        self.overlap_win_size = int(window_size * overlap_ratio) + window_size\n",
    "\n",
    "        self.norm1 = norm_layer(dim)\n",
    "        self.qkv = nn.Linear(dim, dim * 3,  bias=qkv_bias)\n",
    "        self.unfold = nn.Unfold(kernel_size=(self.overlap_win_size, self.overlap_win_size), stride=window_size, padding=(self.overlap_win_size-window_size)//2)\n",
    "\n",
    "        # define a parameter table of relative position bias\n",
    "        self.relative_position_bias_table = nn.Parameter(\n",
    "            torch.zeros((window_size + self.overlap_win_size - 1) * (window_size + self.overlap_win_size - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH\n",
    "\n",
    "        trunc_normal_(self.relative_position_bias_table, std=.02)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "        self.proj = nn.Linear(dim,dim)\n",
    "\n",
    "        self.norm2 = norm_layer(dim)\n",
    "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
    "        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU)\n",
    "\n",
    "    def forward(self, x, x_size, rpi):\n",
    "        h, w = x_size\n",
    "        print(x.shape)\n",
    "        b,c,_,_ = x.shape\n",
    "\n",
    "        shortcut = x\n",
    "        x = self.norm1(x)\n",
    "        x = x.view(b, h, w, c)\n",
    "\n",
    "        qkv = self.qkv(x).reshape(b, h, w, 3, c).permute(3, 0, 4, 1, 2) # 3, b, c, h, w\n",
    "        q = qkv[0].permute(0, 2, 3, 1) # b, h, w, c\n",
    "        kv = torch.cat((qkv[1], qkv[2]), dim=1) # b, 2*c, h, w\n",
    "\n",
    "        # partition windows\n",
    "        q_windows = window_partition(q, self.window_size)  # nw*b, window_size, window_size, c\n",
    "        q_windows = q_windows.view(-1, self.window_size * self.window_size, c)  # nw*b, window_size*window_size, c\n",
    "\n",
    "        kv_windows = self.unfold(kv) # b, c*w*w, nw\n",
    "        kv_windows = rearrange(kv_windows, 'b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch', nc=2, ch=c, owh=self.overlap_win_size, oww=self.overlap_win_size).contiguous() # 2, nw*b, ow*ow, c\n",
    "        k_windows, v_windows = kv_windows[0], kv_windows[1] # nw*b, ow*ow, c\n",
    "\n",
    "        b_, nq, _ = q_windows.shape\n",
    "        _, n, _ = k_windows.shape\n",
    "        d = self.dim // self.num_heads\n",
    "        q = q_windows.reshape(b_, nq, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, nq, d\n",
    "        k = k_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, n, d\n",
    "        v = v_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, n, d\n",
    "\n",
    "        q = q * self.scale\n",
    "        attn = (q @ k.transpose(-2, -1))\n",
    "\n",
    "        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(\n",
    "            self.window_size * self.window_size, self.overlap_win_size * self.overlap_win_size, -1)  # ws*ws, wse*wse, nH\n",
    "        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, ws*ws, wse*wse\n",
    "        attn = attn + relative_position_bias.unsqueeze(0)\n",
    "\n",
    "        attn = self.softmax(attn)\n",
    "        attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim)\n",
    "\n",
    "        # merge windows\n",
    "        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.dim)\n",
    "        x = window_reverse(attn_windows, self.window_size, h, w)  # b h w c\n",
    "        x = x.view(b, h * w, self.dim)\n",
    "\n",
    "        x = self.proj(x) + shortcut\n",
    "\n",
    "        x = x + self.mlp(self.norm2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "4b4e5d21-2b62-45c2-915d-1e6db46fd29b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class OCAB(nn.Module):\n",
    "    # overlapping cross-attention block\n",
    "\n",
    "    def __init__(self, dim,\n",
    "                input_resolution,\n",
    "                window_size,\n",
    "                overlap_ratio,\n",
    "                num_heads,\n",
    "                qkv_bias=True,\n",
    "                qk_scale=None,\n",
    "                mlp_ratio=2,\n",
    "                norm_layer=nn.LayerNorm\n",
    "                ):\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.input_resolution = input_resolution\n",
    "        self.window_size = window_size\n",
    "        self.num_heads = num_heads\n",
    "        head_dim = dim // num_heads\n",
    "        self.scale = qk_scale or head_dim**-0.5\n",
    "        self.overlap_win_size = int(window_size * overlap_ratio) + window_size\n",
    "\n",
    "        self.norm1 = norm_layer(dim)\n",
    "        self.qkv = nn.Linear(dim, dim * 3,  bias=qkv_bias)\n",
    "        self.unfold = nn.Unfold(kernel_size=(self.overlap_win_size, self.overlap_win_size), stride=window_size, padding=(self.overlap_win_size-window_size)//2)\n",
    "\n",
    "        # define a parameter table of relative position bias\n",
    "        self.relative_position_bias_table = nn.Parameter(\n",
    "            torch.zeros((window_size + self.overlap_win_size - 1) * (window_size + self.overlap_win_size - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH\n",
    "\n",
    "        trunc_normal_(self.relative_position_bias_table, std=.02)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "        self.proj = nn.Linear(dim,dim)\n",
    "\n",
    "        self.norm2 = norm_layer(dim)\n",
    "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
    "        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=nn.GELU)\n",
    "\n",
    "    def forward(self, x, x_size, rpi):\n",
    "        h, w = x_size\n",
    "        b, _, c = x.shape\n",
    "\n",
    "        shortcut = x\n",
    "        x = self.norm1(x)\n",
    "        x = x.view(b, h, w, c)\n",
    "\n",
    "        qkv = self.qkv(x).reshape(b, h, w, 3, c).permute(3, 0, 4, 1, 2) # 3, b, c, h, w\n",
    "        q = qkv[0].permute(0, 2, 3, 1) # b, h, w, c\n",
    "        kv = torch.cat((qkv[1], qkv[2]), dim=1) # b, 2*c, h, w\n",
    "\n",
    "        # partition windows\n",
    "        q_windows = window_partition(q, self.window_size)  # nw*b, window_size, window_size, c\n",
    "        q_windows = q_windows.view(-1, self.window_size * self.window_size, c)  # nw*b, window_size*window_size, c\n",
    "\n",
    "        kv_windows = self.unfold(kv) # b, c*w*w, nw\n",
    "        kv_windows = rearrange(kv_windows, 'b (nc ch owh oww) nw -> nc (b nw) (owh oww) ch', nc=2, ch=c, owh=self.overlap_win_size, oww=self.overlap_win_size).contiguous() # 2, nw*b, ow*ow, c\n",
    "        k_windows, v_windows = kv_windows[0], kv_windows[1] # nw*b, ow*ow, c\n",
    "\n",
    "        b_, nq, _ = q_windows.shape\n",
    "        _, n, _ = k_windows.shape\n",
    "        d = self.dim // self.num_heads\n",
    "        q = q_windows.reshape(b_, nq, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, nq, d\n",
    "        k = k_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, n, d\n",
    "        v = v_windows.reshape(b_, n, self.num_heads, d).permute(0, 2, 1, 3) # nw*b, nH, n, d\n",
    "\n",
    "        q = q * self.scale\n",
    "        attn = (q @ k.transpose(-2, -1))\n",
    "\n",
    "        relative_position_bias = self.relative_position_bias_table[rpi.view(-1)].view(\n",
    "            self.window_size * self.window_size, self.overlap_win_size * self.overlap_win_size, -1)  # ws*ws, wse*wse, nH\n",
    "        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, ws*ws, wse*wse\n",
    "        attn = attn + relative_position_bias.unsqueeze(0)\n",
    "\n",
    "        attn = self.softmax(attn)\n",
    "        attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim)\n",
    "\n",
    "        # merge windows\n",
    "        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.dim)\n",
    "        x = window_reverse(attn_windows, self.window_size, h, w)  # b h w c\n",
    "        x = x.view(b, h * w, self.dim)\n",
    "\n",
    "        x = self.proj(x) + shortcut\n",
    "\n",
    "        x = x + self.mlp(self.norm2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "ba2fc8e7-dd1d-4508-b0d9-6e6ba024a93d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttenBlocks(nn.Module):\n",
    "    \"\"\" A series of attention blocks for one RHAG.\n",
    "\n",
    "    Args:\n",
    "        dim (int): Number of input channels.\n",
    "        input_resolution (tuple[int]): Input resolution.\n",
    "        depth (int): Number of blocks.\n",
    "        num_heads (int): Number of attention heads.\n",
    "        window_size (int): Local window size.\n",
    "        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n",
    "        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n",
    "        drop (float, optional): Dropout rate. Default: 0.0\n",
    "        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n",
    "        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n",
    "        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n",
    "        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n",
    "        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 dim,\n",
    "                 input_resolution,\n",
    "                 depth,\n",
    "                 num_heads,\n",
    "                 window_size,\n",
    "                 compress_ratio,\n",
    "                 squeeze_factor,\n",
    "                 conv_scale,\n",
    "                 overlap_ratio,\n",
    "                 mlp_ratio=4.,\n",
    "                 qkv_bias=True,\n",
    "                 qk_scale=None,\n",
    "                 drop=0.,\n",
    "                 attn_drop=0.,\n",
    "                 drop_path=0.,\n",
    "                 norm_layer=nn.LayerNorm,\n",
    "                 downsample=None,\n",
    "                 use_checkpoint=False):\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.input_resolution = input_resolution\n",
    "        self.depth = depth\n",
    "        self.use_checkpoint = use_checkpoint\n",
    "\n",
    "        # build blocks\n",
    "        self.blocks = nn.ModuleList([\n",
    "            HAB(\n",
    "                dim=dim,\n",
    "                input_resolution=input_resolution,\n",
    "                num_heads=num_heads,\n",
    "                window_size=window_size,\n",
    "                shift_size=0 if (i % 2 == 0) else window_size // 2,\n",
    "                compress_ratio=compress_ratio,\n",
    "                squeeze_factor=squeeze_factor,\n",
    "                conv_scale=conv_scale,\n",
    "                mlp_ratio=mlp_ratio,\n",
    "                qkv_bias=qkv_bias,\n",
    "                qk_scale=qk_scale,\n",
    "                drop=drop,\n",
    "                attn_drop=attn_drop,\n",
    "                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n",
    "                norm_layer=norm_layer) for i in range(depth)\n",
    "        ])\n",
    "\n",
    "        # OCAB\n",
    "        self.overlap_attn = OCAB(\n",
    "                            dim=dim,\n",
    "                            input_resolution=input_resolution,\n",
    "                            window_size=window_size,\n",
    "                            overlap_ratio=overlap_ratio,\n",
    "                            num_heads=num_heads,\n",
    "                            qkv_bias=qkv_bias,\n",
    "                            qk_scale=qk_scale,\n",
    "                            mlp_ratio=mlp_ratio,\n",
    "                            norm_layer=norm_layer\n",
    "                            )\n",
    "\n",
    "        # patch merging layer\n",
    "        if downsample is not None:\n",
    "            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)\n",
    "        else:\n",
    "            self.downsample = None\n",
    "\n",
    "    def forward(self, x, x_size, params):\n",
    "        for blk in self.blocks:\n",
    "            x = blk(x, x_size, params['rpi_sa'], params['attn_mask'])\n",
    "\n",
    "        x = self.overlap_attn(x, x_size, params['rpi_oca'])\n",
    "\n",
    "        if self.downsample is not None:\n",
    "            x = self.downsample(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "1e433423-4f27-42c7-95de-8dcf22139963",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RHAG(nn.Module):\n",
    "    \"\"\"Residual Hybrid Attention Group (RHAG).\n",
    "\n",
    "    Args:\n",
    "        dim (int): Number of input channels.\n",
    "        input_resolution (tuple[int]): Input resolution.\n",
    "        depth (int): Number of blocks.\n",
    "        num_heads (int): Number of attention heads.\n",
    "        window_size (int): Local window size.\n",
    "        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n",
    "        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n",
    "        drop (float, optional): Dropout rate. Default: 0.0\n",
    "        attn_drop (float, optional): Attention dropout rate. Default: 0.0\n",
    "        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n",
    "        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n",
    "        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n",
    "        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n",
    "        img_size: Input image size.\n",
    "        patch_size: Patch size.\n",
    "        resi_connection: The convolutional block before residual connection.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 dim,\n",
    "                 input_resolution,\n",
    "                 depth,\n",
    "                 num_heads,\n",
    "                 window_size,\n",
    "                 compress_ratio,\n",
    "                 squeeze_factor,\n",
    "                 conv_scale,\n",
    "                 overlap_ratio,\n",
    "                 mlp_ratio=4.,\n",
    "                 qkv_bias=True,\n",
    "                 qk_scale=None,\n",
    "                 drop=0.,\n",
    "                 attn_drop=0.,\n",
    "                 drop_path=0.,\n",
    "                 norm_layer=nn.LayerNorm,\n",
    "                 downsample=None,\n",
    "                 use_checkpoint=False,\n",
    "                 img_size=256,\n",
    "                 patch_size=4,\n",
    "                 resi_connection='1conv'):\n",
    "        super(RHAG, self).__init__()\n",
    "\n",
    "        self.dim = dim\n",
    "        self.input_resolution = input_resolution\n",
    "\n",
    "        self.residual_group = AttenBlocks(\n",
    "            dim=dim,\n",
    "            input_resolution=input_resolution,\n",
    "            depth=depth,\n",
    "            num_heads=num_heads,\n",
    "            window_size=window_size,\n",
    "            compress_ratio=compress_ratio,\n",
    "            squeeze_factor=squeeze_factor,\n",
    "            conv_scale=conv_scale,\n",
    "            overlap_ratio=overlap_ratio,\n",
    "            mlp_ratio=mlp_ratio,\n",
    "            qkv_bias=qkv_bias,\n",
    "            qk_scale=qk_scale,\n",
    "            drop=drop,\n",
    "            attn_drop=attn_drop,\n",
    "            drop_path=drop_path,\n",
    "            norm_layer=norm_layer,\n",
    "            downsample=downsample,\n",
    "            use_checkpoint=use_checkpoint)\n",
    "\n",
    "        if resi_connection == '1conv':\n",
    "            self.conv = nn.Conv2d(dim, dim, 3, 1, 1)\n",
    "        elif resi_connection == 'identity':\n",
    "            self.conv = nn.Identity()\n",
    "\n",
    "        self.patch_embed = PatchEmbed(\n",
    "            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, norm_layer=None)\n",
    "\n",
    "        self.patch_unembed = PatchUnEmbed(\n",
    "            img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim, norm_layer=None)\n",
    "\n",
    "    def forward(self, x, x_size, params):\n",
    "        return self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size, params), x_size))) + x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "3410f7f4-e718-4bec-9971-666e4cde6538",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "class PatchEmbed(nn.Module):\n",
    "    r\"\"\" Image to Patch Embedding\n",
    "\n",
    "    Args:\n",
    "        img_size (int): Image size.  Default: 224.\n",
    "        patch_size (int): Patch token size. Default: 4.\n",
    "        in_chans (int): Number of input image channels. Default: 3.\n",
    "        embed_dim (int): Number of linear projection output channels. Default: 96.\n",
    "        norm_layer (nn.Module, optional): Normalization layer. Default: None\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, img_size=256, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):\n",
    "        super().__init__()\n",
    "        img_size = to_2tuple(img_size)\n",
    "        patch_size = to_2tuple(patch_size)\n",
    "        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]\n",
    "        self.img_size = img_size\n",
    "        self.patch_size = patch_size\n",
    "        self.patches_resolution = patches_resolution\n",
    "        self.num_patches = patches_resolution[0] * patches_resolution[1]\n",
    "\n",
    "        self.in_chans = in_chans\n",
    "        self.embed_dim = embed_dim\n",
    "\n",
    "        if norm_layer is not None:\n",
    "            self.norm = norm_layer(embed_dim)\n",
    "        else:\n",
    "            self.norm = None\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.flatten(2).transpose(1, 2)  # b Ph*Pw c\n",
    "        if self.norm is not None:\n",
    "            x = self.norm(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class PatchUnEmbed(nn.Module):\n",
    "    r\"\"\" Image to Patch Unembedding\n",
    "\n",
    "    Args:\n",
    "        img_size (int): Image size.  Default: 224.\n",
    "        patch_size (int): Patch token size. Default: 4.\n",
    "        in_chans (int): Number of input image channels. Default: 3.\n",
    "        embed_dim (int): Number of linear projection output channels. Default: 96.\n",
    "        norm_layer (nn.Module, optional): Normalization layer. Default: None\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):\n",
    "        super().__init__()\n",
    "        img_size = to_2tuple(img_size)\n",
    "        patch_size = to_2tuple(patch_size)\n",
    "        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]\n",
    "        self.img_size = img_size\n",
    "        self.patch_size = patch_size\n",
    "        self.patches_resolution = patches_resolution\n",
    "        self.num_patches = patches_resolution[0] * patches_resolution[1]\n",
    "\n",
    "        self.in_chans = in_chans\n",
    "        self.embed_dim = embed_dim\n",
    "\n",
    "    def forward(self, x, x_size):\n",
    "        x = x.transpose(1, 2).contiguous().view(x.shape[0], self.embed_dim, x_size[0], x_size[1])  # b Ph*Pw c\n",
    "        return x\n",
    "\n",
    "\n",
    "class Upsample(nn.Sequential):\n",
    "    \"\"\"Upsample module.\n",
    "\n",
    "    Args:\n",
    "        scale (int): Scale factor. Supported scales: 2^n and 3.\n",
    "        num_feat (int): Channel number of intermediate features.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, scale, num_feat):\n",
    "        m = []\n",
    "        if (scale & (scale - 1)) == 0:  # scale = 2^n\n",
    "            for _ in range(int(math.log(scale, 2))):\n",
    "                m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))\n",
    "                m.append(nn.PixelShuffle(2))\n",
    "        elif scale == 3:\n",
    "            m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))\n",
    "            m.append(nn.PixelShuffle(3))\n",
    "        else:\n",
    "            raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')\n",
    "        super(Upsample, self).__init__(*m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "a5cf795f-7f15-4aa3-900f-70adb3907a76",
   "metadata": {},
   "outputs": [],
   "source": [
    "class HAT(nn.Module):\n",
    "    r\"\"\" Hybrid Attention Transformer\n",
    "        A PyTorch implementation of : `Activating More Pixels in Image Super-Resolution Transformer`.\n",
    "        Some codes are based on SwinIR.\n",
    "    Args:\n",
    "        img_size (int | tuple(int)): Input image size. Default 64\n",
    "        patch_size (int | tuple(int)): Patch size. Default: 1\n",
    "        in_chans (int): Number of input image channels. Default: 3\n",
    "        embed_dim (int): Patch embedding dimension. Default: 96\n",
    "        depths (tuple(int)): Depth of each Swin Transformer layer.\n",
    "        num_heads (tuple(int)): Number of attention heads in different layers.\n",
    "        window_size (int): Window size. Default: 7\n",
    "        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4\n",
    "        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True\n",
    "        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None\n",
    "        drop_rate (float): Dropout rate. Default: 0\n",
    "        attn_drop_rate (float): Attention dropout rate. Default: 0\n",
    "        drop_path_rate (float): Stochastic depth rate. Default: 0.1\n",
    "        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.\n",
    "        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False\n",
    "        patch_norm (bool): If True, add normalization after patch embedding. Default: True\n",
    "        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False\n",
    "        upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction\n",
    "        img_range: Image range. 1. or 255.\n",
    "        upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None\n",
    "        resi_connection: The convolutional block before residual connection. '1conv'/'3conv'\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 img_size=256,\n",
    "                 patch_size=1,\n",
    "                 in_chans=3,\n",
    "                 embed_dim=96,\n",
    "                 depths=(6, 6, 6, 6),\n",
    "                 num_heads=(6, 6, 6, 6),\n",
    "                 window_size=8,\n",
    "                 compress_ratio=3,\n",
    "                 squeeze_factor=30,\n",
    "                 conv_scale=0.01,\n",
    "                 overlap_ratio=0.5,\n",
    "                 mlp_ratio=4.,\n",
    "                 qkv_bias=True,\n",
    "                 qk_scale=None,\n",
    "                 drop_rate=0.1,\n",
    "                 attn_drop_rate=0.,\n",
    "                 drop_path_rate=0.1,\n",
    "                 norm_layer=nn.LayerNorm,\n",
    "                 ape=False,\n",
    "                 patch_norm=True,\n",
    "                 use_checkpoint=False,\n",
    "                 upscale=2,\n",
    "                 img_range=1.,\n",
    "                 upsampler='pixelshuffle',\n",
    "                 resi_connection='1conv',\n",
    "                 **kwargs):\n",
    "        super(HAT, self).__init__()\n",
    "\n",
    "        self.window_size = window_size\n",
    "        self.shift_size = window_size // 2\n",
    "        self.overlap_ratio = overlap_ratio\n",
    "\n",
    "        num_in_ch = in_chans\n",
    "        num_out_ch = in_chans\n",
    "        num_feat = 64\n",
    "        self.img_range = img_range\n",
    "        if in_chans == 3:\n",
    "            rgb_mean = (0.4488, 0.4371, 0.4040)\n",
    "            self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)\n",
    "        else:\n",
    "            self.mean = torch.zeros(1, 1, 1, 1)\n",
    "        self.upscale = upscale\n",
    "        self.upsampler = upsampler\n",
    "\n",
    "        # relative position index\n",
    "        relative_position_index_SA = self.calculate_rpi_sa()\n",
    "        relative_position_index_OCA = self.calculate_rpi_oca()\n",
    "        self.register_buffer('relative_position_index_SA', relative_position_index_SA)\n",
    "        self.register_buffer('relative_position_index_OCA', relative_position_index_OCA)\n",
    "\n",
    "        # ------------------------- 1, shallow feature extraction ------------------------- #\n",
    "        self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)\n",
    "\n",
    "        # ------------------------- 2, deep feature extraction ------------------------- #\n",
    "        self.num_layers = len(depths)\n",
    "        self.embed_dim = embed_dim\n",
    "        self.ape = ape\n",
    "        self.patch_norm = patch_norm\n",
    "        self.num_features = embed_dim\n",
    "        self.mlp_ratio = mlp_ratio\n",
    "\n",
    "        # split image into non-overlapping patches\n",
    "        self.patch_embed = PatchEmbed(\n",
    "            img_size=img_size,\n",
    "            patch_size=patch_size,\n",
    "            in_chans=embed_dim,\n",
    "            embed_dim=embed_dim,\n",
    "            norm_layer=norm_layer if self.patch_norm else None)\n",
    "        num_patches = self.patch_embed.num_patches\n",
    "        patches_resolution = self.patch_embed.patches_resolution\n",
    "        self.patches_resolution = patches_resolution\n",
    "\n",
    "        # merge non-overlapping patches into image\n",
    "        self.patch_unembed = PatchUnEmbed(\n",
    "            img_size=img_size,\n",
    "            patch_size=patch_size,\n",
    "            in_chans=embed_dim,\n",
    "            embed_dim=embed_dim,\n",
    "            norm_layer=norm_layer if self.patch_norm else None)\n",
    "\n",
    "        # absolute position embedding\n",
    "        if self.ape:\n",
    "            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))\n",
    "            trunc_normal_(self.absolute_pos_embed, std=.02)\n",
    "\n",
    "        self.pos_drop = nn.Dropout(p=drop_rate)\n",
    "\n",
    "        # stochastic depth\n",
    "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule\n",
    "\n",
    "        # build Residual Hybrid Attention Groups (RHAG)\n",
    "        self.layers = nn.ModuleList()\n",
    "        for i_layer in range(self.num_layers):\n",
    "            layer = RHAG(\n",
    "                dim=embed_dim,\n",
    "                input_resolution=(patches_resolution[0], patches_resolution[1]),\n",
    "                depth=depths[i_layer],\n",
    "                num_heads=num_heads[i_layer],\n",
    "                window_size=window_size,\n",
    "                compress_ratio=compress_ratio,\n",
    "                squeeze_factor=squeeze_factor,\n",
    "                conv_scale=conv_scale,\n",
    "                overlap_ratio=overlap_ratio,\n",
    "                mlp_ratio=self.mlp_ratio,\n",
    "                qkv_bias=qkv_bias,\n",
    "                qk_scale=qk_scale,\n",
    "                drop=drop_rate,\n",
    "                attn_drop=attn_drop_rate,\n",
    "                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],  # no impact on SR results\n",
    "                norm_layer=norm_layer,\n",
    "                downsample=None,\n",
    "                use_checkpoint=use_checkpoint,\n",
    "                img_size=img_size,\n",
    "                patch_size=patch_size,\n",
    "                resi_connection=resi_connection)\n",
    "            self.layers.append(layer)\n",
    "        self.norm = norm_layer(self.num_features)\n",
    "\n",
    "        # build the last conv layer in deep feature extraction\n",
    "        if resi_connection == '1conv':\n",
    "            self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)\n",
    "        elif resi_connection == 'identity':\n",
    "            self.conv_after_body = nn.Identity()\n",
    "\n",
    "        # ------------------------- 3, high quality image reconstruction ------------------------- #\n",
    "        if self.upsampler == 'pixelshuffle':\n",
    "            # for classical SR\n",
    "            self.conv_before_upsample = nn.Sequential(\n",
    "                nn.Conv2d(embed_dim, num_feat, 3, 1, 1), nn.LeakyReLU(inplace=True))\n",
    "            self.upsample = Upsample(upscale, num_feat)\n",
    "            self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)\n",
    "\n",
    "        self.apply(self._init_weights)\n",
    "\n",
    "    def _init_weights(self, m):\n",
    "        if isinstance(m, nn.Linear):\n",
    "            trunc_normal_(m.weight, std=.02)\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "        elif isinstance(m, nn.LayerNorm):\n",
    "            nn.init.constant_(m.bias, 0)\n",
    "            nn.init.constant_(m.weight, 1.0)\n",
    "\n",
    "    def calculate_rpi_sa(self):\n",
    "        # calculate relative position index for SA\n",
    "        coords_h = torch.arange(self.window_size)\n",
    "        coords_w = torch.arange(self.window_size)\n",
    "        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww\n",
    "        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww\n",
    "        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww\n",
    "        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2\n",
    "        relative_coords[:, :, 0] += self.window_size - 1  # shift to start from 0\n",
    "        relative_coords[:, :, 1] += self.window_size - 1\n",
    "        relative_coords[:, :, 0] *= 2 * self.window_size - 1\n",
    "        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww\n",
    "        return relative_position_index\n",
    "\n",
    "    def calculate_rpi_oca(self):\n",
    "        # calculate relative position index for OCA\n",
    "        window_size_ori = self.window_size\n",
    "        window_size_ext = self.window_size + int(self.overlap_ratio * self.window_size)\n",
    "\n",
    "        coords_h = torch.arange(window_size_ori)\n",
    "        coords_w = torch.arange(window_size_ori)\n",
    "        coords_ori = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, ws, ws\n",
    "        coords_ori_flatten = torch.flatten(coords_ori, 1)  # 2, ws*ws\n",
    "\n",
    "        coords_h = torch.arange(window_size_ext)\n",
    "        coords_w = torch.arange(window_size_ext)\n",
    "        coords_ext = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, wse, wse\n",
    "        coords_ext_flatten = torch.flatten(coords_ext, 1)  # 2, wse*wse\n",
    "\n",
    "        relative_coords = coords_ext_flatten[:, None, :] - coords_ori_flatten[:, :, None]   # 2, ws*ws, wse*wse\n",
    "\n",
    "        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # ws*ws, wse*wse, 2\n",
    "        relative_coords[:, :, 0] += window_size_ori - window_size_ext + 1  # shift to start from 0\n",
    "        relative_coords[:, :, 1] += window_size_ori - window_size_ext + 1\n",
    "\n",
    "        relative_coords[:, :, 0] *= window_size_ori + window_size_ext - 1\n",
    "        relative_position_index = relative_coords.sum(-1)\n",
    "        return relative_position_index\n",
    "\n",
    "    def calculate_mask(self, x_size):\n",
    "        # calculate attention mask for SW-MSA\n",
    "        h, w = x_size\n",
    "        img_mask = torch.zeros((1, h, w, 1))  # 1 h w 1\n",
    "        h_slices = (slice(0, -self.window_size), slice(-self.window_size,\n",
    "                                                       -self.shift_size), slice(-self.shift_size, None))\n",
    "        w_slices = (slice(0, -self.window_size), slice(-self.window_size,\n",
    "                                                       -self.shift_size), slice(-self.shift_size, None))\n",
    "        cnt = 0\n",
    "        for h in h_slices:\n",
    "            for w in w_slices:\n",
    "                img_mask[:, h, w, :] = cnt\n",
    "                cnt += 1\n",
    "\n",
    "        mask_windows = window_partition(img_mask, self.window_size)  # nw, window_size, window_size, 1\n",
    "        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)\n",
    "        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)\n",
    "        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))\n",
    "\n",
    "        return attn_mask\n",
    "\n",
    "    @torch.jit.ignore\n",
    "    def no_weight_decay(self):\n",
    "        return {'absolute_pos_embed'}\n",
    "\n",
    "    @torch.jit.ignore\n",
    "    def no_weight_decay_keywords(self):\n",
    "        return {'relative_position_bias_table'}\n",
    "\n",
    "    def forward_features(self, x):\n",
    "        x_size = (x.shape[2], x.shape[3])\n",
    "\n",
    "        # Calculate attention mask and relative position index in advance to speed up inference. \n",
    "        # The original code is very time-consuming for large window size.\n",
    "        attn_mask = self.calculate_mask(x_size).to(x.device)\n",
    "        params = {'attn_mask': attn_mask, 'rpi_sa': self.relative_position_index_SA, 'rpi_oca': self.relative_position_index_OCA}\n",
    "\n",
    "        x = self.patch_embed(x)\n",
    "        if self.ape:\n",
    "            x = x + self.absolute_pos_embed\n",
    "        x = self.pos_drop(x)\n",
    "\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, x_size, params)\n",
    "\n",
    "        x = self.norm(x)  # b seq_len c\n",
    "        x = self.patch_unembed(x, x_size)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def forward(self, x):\n",
    "        self.mean = self.mean.type_as(x)\n",
    "        x = (x - self.mean) * self.img_range\n",
    "\n",
    "        # if self.upsampler == 'pixelshuffle':\n",
    "        #     # for classical SR\n",
    "        #     x = self.conv_first(x)\n",
    "        #     x = self.conv_after_body(self.forward_features(x)) + x\n",
    "        #     x = self.conv_before_upsample(x)\n",
    "        #     x = self.conv_last(self.upsample(x))\n",
    "        if True:\n",
    "            # for classical SR\n",
    "            x = self.conv_first(x)\n",
    "            x = self.conv_after_body(self.forward_features(x)) + x\n",
    "            x = self.conv_before_upsample(x)\n",
    "            x = self.conv_last(self.upsample(x))\n",
    "        print(x.shape)\n",
    "        x = x / self.img_range + self.mean\n",
    "        print(x.shape)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "eb7e745a-9433-4b8c-b792-5bb56d78f9cb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "models=HAT()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "id": "8c24ceaf-7692-4a0f-91af-f65869a78eec",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 3, 512, 512])\n",
      "torch.Size([1, 3, 512, 512])\n"
     ]
    }
   ],
   "source": [
    "s=models(torch.zeros(1,3,256,256))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "260188c8-dc04-4696-b364-8bb279135a50",
   "metadata": {},
   "outputs": [],
   "source": [
    "from basicsr.utils.registry import ARCH_REGISTRY\n",
    "from torch import nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.nn.utils import spectral_norm\n",
    "\n",
    "\n",
    "@ARCH_REGISTRY.register()\n",
    "class UNetDiscriminatorSN(nn.Module):\n",
    "    \"\"\"Defines a U-Net discriminator with spectral normalization (SN)\n",
    "\n",
    "    It is used in Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.\n",
    "\n",
    "    Arg:\n",
    "        num_in_ch (int): Channel number of inputs. Default: 3.\n",
    "        num_feat (int): Channel number of base intermediate features. Default: 64.\n",
    "        skip_connection (bool): Whether to use skip connections between U-Net. Default: True.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, num_in_ch, num_feat=64, skip_connection=True):\n",
    "        super(UNetDiscriminatorSN, self).__init__()\n",
    "        self.skip_connection = skip_connection\n",
    "        norm = spectral_norm\n",
    "        # the first convolution\n",
    "        self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)\n",
    "        # downsample\n",
    "        self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 4, 2, 1, bias=False))\n",
    "        self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 4, 2, 1, bias=False))\n",
    "        self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 4, 2, 1, bias=False))\n",
    "        # upsample\n",
    "        self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))\n",
    "        self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))\n",
    "        self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))\n",
    "        # extra convolutions\n",
    "        self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))\n",
    "        self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))\n",
    "        self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # downsample\n",
    "        x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)\n",
    "        x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)\n",
    "        x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)\n",
    "        x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)\n",
    "\n",
    "        # upsample\n",
    "        x3 = F.interpolate(x3, scale_factor=2, mode='bilinear', align_corners=False)\n",
    "        x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)\n",
    "\n",
    "        if self.skip_connection:\n",
    "            x4 = x4 + x2\n",
    "        x4 = F.interpolate(x4, scale_factor=2, mode='bilinear', align_corners=False)\n",
    "        x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)\n",
    "\n",
    "        if self.skip_connection:\n",
    "            x5 = x5 + x1\n",
    "        x5 = F.interpolate(x5, scale_factor=2, mode='bilinear', align_corners=False)\n",
    "        x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)\n",
    "\n",
    "        if self.skip_connection:\n",
    "            x6 = x6 + x0\n",
    "\n",
    "        # extra convolutions\n",
    "        out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)\n",
    "        out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)\n",
    "        out = self.conv9(out)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "966f70d4-f716-465d-9941-6452f49dd9d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dis=UNetDiscriminatorSN(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "id": "e1b0fe65-3519-4a7d-81e0-4a3688ecf784",
   "metadata": {},
   "outputs": [],
   "source": [
    "s=dis(torch.zeros(1,3,512,512))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "62df1150-e4d6-4777-9b44-b507b2c0a196",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 512, 512])"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78aac80a-8638-476d-b5ec-eea3b8b15561",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
